"""
分析服务模块
"""
from typing import Dict, Any, Optional
import numpy as np
from sklearn.cluster import KMeans
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd
import warnings

from app.models.analysis import AnalysisTask
from app.models.dataset import Dataset
from app import db

# 忽略特定警告
warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib')
warnings.filterwarnings('ignore', category=FutureWarning, module='sklearn')

class AnalysisService:
    """分析服务类"""
    
    @staticmethod
    def create_task(dataset_id: int, task_type: str, parameters: Dict[str, Any]) -> AnalysisTask:
        """
        创建分析任务
        
        Args:
            dataset_id: 数据集ID
            task_type: 分析类型
            parameters: 分析参数
            
        Returns:
            AnalysisTask: 创建的任务
        """
        if not AnalysisTask.validate_task_type(task_type):
            raise ValueError(f"Invalid task type: {task_type}")
            
        task = AnalysisTask(
            dataset_id=dataset_id,
            task_type=task_type,
            parameters=parameters
        )
        
        db.session.add(task)
        db.session.commit()
        
        return task
        
    @staticmethod
    def get_task(task_id: int) -> Optional[AnalysisTask]:
        """
        获取分析任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            Optional[AnalysisTask]: 任务对象，如果不存在则返回None
        """
        return AnalysisTask.query.get(task_id)
        
    @staticmethod
    def execute_task(task: AnalysisTask):
        """
        执行分析任务
        
        Args:
            task: 分析任务对象
        """
        try:
            task.start()
            
            # 获取数据集
            dataset = Dataset.query.get(task.dataset_id)
            if not dataset:
                raise ValueError(f"Dataset not found: {task.dataset_id}")
                
            # 获取数据
            data = dataset.get_data()  # 假设Dataset模型有get_data方法
            
            # 根据任务类型执行不同的分析
            if task.task_type == AnalysisTask.TYPE_KMEANS:
                result = AnalysisService._execute_kmeans(data, task.parameters)
            elif task.task_type == AnalysisTask.TYPE_LINEAR_REGRESSION:
                result = AnalysisService._execute_linear_regression(data, task.parameters)
            else:
                raise ValueError(f"Unsupported task type: {task.task_type}")
                
            task.complete(result)
            
        except Exception as e:
            task.fail(str(e))
            raise
            
    @staticmethod
    def _execute_kmeans(data: np.ndarray, parameters: Dict) -> Dict:
        """执行K-means聚类"""
        n_clusters = parameters.get('n_clusters', 3)
        
        # 数据标准化
        scaler = StandardScaler()
        scaled_data = scaler.fit_transform(data)
        
        # 执行聚类
        kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
        labels = kmeans.fit_predict(scaled_data)
        
        return {
            'labels': labels.tolist(),
            'centroids': kmeans.cluster_centers_.tolist(),
            'inertia': float(kmeans.inertia_),
            'scaled_data': scaled_data.tolist()
        }
        
    @staticmethod
    def _execute_linear_regression(data: np.ndarray, parameters: Dict) -> Dict:
        """
        执行线性回归
        
        Args:
            data: 输入数据，最后一列为目标变量
            parameters: 参数字典，可包含：
                - test_size: 测试集比例
        """
        # 获取参数
        test_size = parameters.get('test_size', 0.2)
        target_index = parameters.get('target_index', -1)
        X = data[:, :target_index] if target_index != -1 else data[:, :-1]
        y = data[:, target_index] if target_index != -1 else data[:, -1]
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42
        )
        
        # 数据标准化
        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)
        
        # 训练模型
        model = LinearRegression()
        model.fit(X_train_scaled, y_train)
        
        # 预测
        y_pred = model.predict(X_test_scaled)
        
        return {
            'coefficients': model.coef_.tolist(),
            'intercept': float(model.intercept_),
            'r2_score': float(r2_score(y_test, y_pred)),
            'mse': float(mean_squared_error(y_test, y_pred)),
            'predictions': y_pred.tolist(),
            'actual_values': y_test.tolist()
        } 